from game.game import SnakeGame
from train import train
from model.models import RLExpert
from model.replay_buffer import ReplayBuffer
import torch.optim as optim
import torch.nn as nn


def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)


def main():
    env = SnakeGame()

    expert = RLExpert()  # 只使用单个专家
    expert.apply(init_weights)

    buffer = ReplayBuffer()
    optimizer = optim.AdamW(expert.parameters(), lr=1e-4, weight_decay=1e-5)

    train(expert, env, buffer, optimizer, episodes=2000)


if __name__ == "__main__":
    main()
